package com.tianji.aigc.memory;

import cn.hutool.core.collection.CollStreamUtil;
import cn.hutool.core.util.ObjectUtil;
import com.baomidou.mybatisplus.core.toolkit.Wrappers;
import com.tianji.aigc.domain.po.ChatRecord;
import com.tianji.aigc.service.IChatRecordService;
import com.tianji.common.utils.UserContext;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.messages.Message;
import org.springframework.beans.factory.annotation.Autowired;

import java.util.List;
import java.util.stream.Collectors;

/**
 * 4、自定义Mysql会话记忆 - Day1 - 作业3
 */
public class MySQLChatMemory implements ChatMemory {
    @Autowired
    private IChatRecordService chatRecordService;

    /**
     * 1、添加会话记忆
     *
     * @param conversationId 会话ID
     * @param messages       消息集合
     */
    @Override
    public void add(String conversationId, List<Message> messages) {
        // 1、消息集合判空
        if (ObjectUtil.isEmpty(messages)) {
            return;
        }
        // 2、根据消息集合，组装ChatRecord集合
        List<ChatRecord> chatRecords = messages.stream().map(message ->
                ChatRecord.builder()
                        .userId(UserContext.getUser())
                        .sessionId(conversationId)
                        .record(ChatMessage.toJson(message))
                        .build()
        ).collect(Collectors.toList());

        // 3、保存
        chatRecordService.saveBatch(chatRecords);
    }

    /**
     * 2、获取会话记忆
     *
     * @param conversationId 会话ID
     * @param lastN          获取条数
     * @return
     */
    @Override
    public List<Message> get(String conversationId, int lastN) {
        // 1、判断角标合理性
        if (lastN < 0) {
            return List.of();
        }

        // 2、获取会话记忆
        List<String> list = this.chatRecordService.lambdaQuery()
                .eq(ChatRecord::getSessionId, conversationId)
                .eq(ChatRecord::getUserId, UserContext.getUser())
                .orderBy(true, true, ChatRecord::getId)
                .last("limit " + lastN)
                .list()
                .stream()
                .map(ChatRecord::getRecord)
                .collect(Collectors.toList());


        // 3、转换为Message集合
        // return list.stream().map(json -> ChatMessage.toMessage(json)).collect(Collectors.toList());
        // return CollStreamUtil.toList(list, json -> ChatMessage.toMessage(json));
        return CollStreamUtil.toList(list,ChatMessage::toMessage);
    }

    /**
     * 3、清空会话记忆
     *
     * @param conversationId
     */
    @Override
    public void clear(String conversationId) {
        this.chatRecordService.remove(
                Wrappers.<ChatRecord>lambdaQuery()
                        .eq(ChatRecord::getSessionId, conversationId)
                        .eq(ChatRecord::getUserId, UserContext.getUser())
        );
    }
}
